.globl matmul

.text
# =======================================================
# FUNCTION: Matrix Multiplication of 2 integer matrices
#   d = matmul(m0, m1)
# Arguments:
#   a0 (int*)  is the pointer to the start of m0
#   a1 (int)   is the # of rows (height) of m0
#   a2 (int)   is the # of columns (width) of m0
#   a3 (int*)  is the pointer to the start of m1
#   a4 (int)   is the # of rows (height) of m1
#   a5 (int)   is the # of columns (width) of m1
#   a6 (int*)  is the pointer to the the start of d
# Returns:
#   None (void), sets d = matmul(m0, m1)
# Exceptions:
#   Make sure to check in top to bottom order!
#   - If the dimensions of m0 do not make sense,
#     this function terminates the program with exit code 38
#   - If the dimensions of m1 do not make sense,
#     this function terminates the program with exit code 38
#   - If the dimensions of m0 and m1 don't match,
#     this function terminates the program with exit code 38
# =======================================================
matmul:
	# Error checks
    addi t0, x0, 1
    blt a1, t0, error
    blt a2, t0, error
    blt a4, t0, error
    blt a5, t0, error
    bne a2, a4, error
	# Prologue
    mv t0, a0 # address of m0
    mv t2, x0 # i = 0
    mv t6, a2
    slli t6, t6, 2
outer_loop_start:
    beq t2, a1, outer_loop_end
    mv t4, x0   # j = 0
    mv a0, t0   # address of m0 is in t0
    mv t1, a3   # address of m1
inner_loop_start:
    beq t4, a5, inner_loop_end
    # Save Argument Registers
    addi sp, sp, -52
    sw t1, 0(sp)
    sw t2, 4(sp)
    sw t3, 8(sp)
    sw t4, 12(sp)
    sw a0, 16(sp)
    sw a1, 20(sp)
    sw a2, 24(sp)
    sw a5, 28(sp)
    sw a6, 32(sp)
    sw t6, 36(sp)
    sw t0, 40(sp)
    sw ra, 44(sp)
    sw a3, 48(sp)
    mv a1, t1   # address of m1 is in t1
    li a3, 1    # stride of arr0
    mv a4, a5   # stride of arr1
    # call dot
    jal ra, dot
    lw a6, 32(sp)
    # Return value should store in d
    sw a0, 0(a6)
    # Recover after store
    lw t1, 0(sp)
    lw t2, 4(sp)
    lw t3, 8(sp)
    lw t4, 12(sp)
    lw a0, 16(sp)
    lw a1, 20(sp)
    lw a2, 24(sp)
    lw a5, 28(sp)
    lw t6, 36(sp)
    lw t0, 40(sp)
    lw ra, 44(sp)
    lw a3, 48(sp)
    addi a6, a6, 4
    addi sp, sp, 52
    # update address of m1
    addi t1, t1, 4  # point the next column of m1
    addi t4, t4, 1  # j++
    j inner_loop_start
inner_loop_end:
    # update address of m0
    add t0, t0, t6  # point the next row of m0
    addi t2, t2, 1  # i++
    j outer_loop_start
outer_loop_end:
	# Epilogue
	ret
error:
    li a0, 38
    j exit